#!/usr/bin/env python
# vim: set fileencoding=utf-8 :
#
# This file is part of Coruja-scripts
#
# Coruja-scripts is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, version 3 of the License.
#
# Coruja-scripts is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Coruja-scripts.  If not, see <http://www.gnu.org/licenses/>.
#
# Copyright 2011 Grupo Falabrasil - http://www.laps.ufpa.br/falabrasil
#
# Author 2011: Pedro Batista - pedosb@gmail.com

import argparse
import sys
import random
import os
import errno
import scipy
from scipy.io import wavfile
from scipy import signal
from scipy.signal import filter_design as fd
import numpy as np
try:
	from matplotlib import pyplot as plt
except ImportError:
	plt = None

from HTKParameters import HTKParameters as hpm
from config import *

class Signal():
	def __init__(self, wav_fn, window_type=None, window_len=25, frame_rate=100):
		"""
		window_type = None (hamming)
		frame_rate => Hz
		window_len => msec
		"""
		try:
			self.wav_fs, self.wav_data = wavfile.read(wav_fn)
		except:
			print wav_fn
			raise
		self.frame_rate = frame_rate
		self.frame_period = self.wav_fs / frame_rate
		self.window_len = window_len * (1./1000)
		self.window = self.get_window(window_type)
		self.make_fbank()
#		self.q1 = signal.firwin(100, [0.001, 50], pass_zero=False, nyq=8000)
#		self.q2 = signal.firwin(100, [50, 100], pass_zero=False, nyq=8000)

	def make_fbank(self, plot=False):
		fir_len = 300
		C = 150
		alfa = 4./3.5
#		alfa = 2
		b = [C]
		f = [100]
		self.q = [signal.firwin(fir_len, [f[0] - b[0] / 2, f[0] + b[0] / 2],
			pass_zero=False, nyq=self.wav_fs/2, window=('hamming'))]
#		for i in xrange(1, Q):
		i = 0
		while True:
			i += 1
			b.append(alfa * b[i-1])
			f.append(f[0] + sum(b[:i]) + (b[i] - b[0]) / 2)
			try:
				self.q.append(signal.firwin(fir_len, [f[i] - b[i] / 2, f[i] + b[i] / 2],
					pass_zero=False, nyq=self.wav_fs/2, window=('hamming')))
			except ValueError:
				break
		self.f = f[:len(f)-1]
#		print self.f
#		print b
		if plot:
			print_sum = True
			if print_sum:
				plt.plot(signal.fftfreq(10000, 1./16000), 10 * np.log10(abs(signal.fft(sum(self.q), 10000))))
				plt.figure()
			for qi in self.q:
				plt.plot(signal.fftfreq(10000, 1./16000), 10 * np.log10( abs(signal.fft(qi, 10000))))
				plt.hold(True)
#			plt.plot([f[i] - b[i] / 2 for i in xrange(len(f))], [1] * len(f), 'o')
#			plt.plot([f[i] + b[i] / 2 for i in xrange(len(f))], [1] * len(f), 'o')
			plt.show()

	def get_window(self, window_type):
		N = self.window_len * self.wav_fs
		return np.hamming(N)
#		return 0.54-0.46*np.cos(np.arange(N)*2*np.pi/(N-1))
#		return np.array([1] * 200)

	def _calc_energy(self, n):
		m_min = n - len(self.window) + 1
		#print 'm_min',m_min
		m_max = m_min + len(self.window)
		if m_min < 0:
			m_min = 0
		if m_max > len(self.wav_data):
			m_max = len(self.wav_data)
		#print 'm_min',m_min
		#print 'm_max',m_max
		q = 0
		for m in np.arange(m_min, m_max):
			q += np.power(self.wav_data[m],2) * self.window[n-m]
		return np.log(q)
#	def _calc_energy(self, n):
#		if n > len(self.wav_data) - len(self.window):
#			return 0
#		if n < 0:
#			return 0
#		return np.log(np.sum(np.power(self.wav_data[n:n+len(self.window)], 2) * self.window))

	def get_frame(self, n):
		if n > len(self.wav_data) - len(self.window):
			return 0
		elif n < 0:
			return 0
		return self.wav_data[n:n+len(self.window)] * self.window

	def _calc_fbank(self, frame=None):
		s = []
		frame_value = self.get_frame(frame)
		for qi in self.q:
			s.append(10 * np.log10(sum([np.power(np.abs(i), 2) for i in signal.lfilter(qi, [1], frame_value)])))
		return s

	def get_energy(self, n=None):
		if n is None:
			energy = []
			for i in xrange(0, len(self.wav_data) + len(self.window), self.frame_period):
				energy.append(self._calc_energy(i))
			return energy
#			return [self._calc_energy(i) for i in xrange(len(self.wav_data))]
		return self._calc_energy(n)

	def get_fbank(self, frame=None, progress=True, plot=False):
		if frame is None:
			frame = range(0, len(self.wav_data)/self.frame_period)
			while True:
				if (frame[len(frame)-1] * self.frame_period + len(self.window)) > len(self.wav_data):
					frame = frame[:len(frame)-1]
				else:
					break
		ft = []
		for i in frame:
			if progress:
				sys.stdout.write("\r%2.2f" % ((float(i) / (len(frame) - 1)) * 100))
				sys.stdout.flush()
			ft.append(self._calc_fbank(i * self.frame_period))
		if progress:
			print
		if plot:
			fta = np.array(ft)
			conf = ''
			for i in xrange(fta.shape[1]):
				if i % 6 == 0 and i != 0:
					conf = '--'
				plt.plot(fta[:,i], conf)
				plt.hold(True)
			plt.legend([str(i) for i in self.f])
			plt.show()
		return ft

def get_htk_parameters(wav_file):
	s = Signal(wav_file, window_len=25, frame_rate=100)
	return hpm(samples=s.get_fbank(plot=False),
			samp_period=s.frame_period)

def split_list_mpi(my_list):
	from mpi4py import MPI
	comm = MPI.COMM_WORLD
	size = MPI.COMM_WORLD.Get_size()
	rank = MPI.COMM_WORLD.Get_rank()
	l = my_list
	if rank == 0:
		random.shuffle(l)
	else:
		l = None
	l = comm.bcast(l, root=0)
	m = len(l) / size
	resto = len(l) - m * size
	s = rank * m + (rank if rank < resto else resto)
	l = l[s : s + m + (1 if rank < resto else 0)]
	return l

def mkdir_p(path):
	"""
	http://stackoverflow.com/questions/600268/mkdir-p-functionality-in-python
	"""
	try:
		os.makedirs(path)
	except OSError as exc: # Python >2.5
		if exc.errno == errno.EEXIST:
			pass
		else: raise

def front_end_mpi(file_list, mask, replacement):
	import re
	for f in split_list_mpi([i.strip() for i in open(file_list)]):
		d = re.sub(mask, replacement, f)
		run_front_end(f, d)

def run_front_end(input_file, output_file):
	s = Signal(input_file, window_len=25, frame_rate=100)
	d = os.path.dirname(output_file)
	if d:
		mkdir_p(d)
	f = np.array(s.get_fbank(plot=False))
	#d = scipy.zeros(f.shape)
	#dd = scipy.zeros(f.shape)
	#for i in xrange(f.shape[1]):
	#	d[:,i] = derivative(f[:,i], 1)
	#	dd[:,i] = derivative(f[:,i], 2)
	#f = scipy.concatenate((f,d,dd), axis=1)
	hpm(output_file, f, s.frame_period)

def derivative(l, n):
	size = len(l)
	ml = scipy.append(l, l[size-5:size-1])
	ml = scipy.append(ml[0:4], ml)
	return np.array([scipy.derivative(ml.__getitem__, i, 1, n, order=5)
			for i in xrange(4,size+4)])

if __name__=='__main__':
	parser = argparse.ArgumentParser(description='Calculate the FBANK front end')
	parser.add_argument('-f', '--input-file',
			help='WAV input file name')
	parser.add_argument('-o', '--output-file',
			help='HPM output file name')
	parser.add_argument('-l', '--input-file-list',
			help='List of WAV input files name')
	parser.add_argument('-p', '--plot-filter', action='store_true',
			help='Print the filterbank')
	parser.add_argument('-m', '--input-file-mask',
			default=config.get('files', 'front_end_input_mask'),
			help='Mask to match the input file')
	parser.add_argument('-r', '--replacement',
			default=config.get('files', 'front_end_input_replacement'),
			help='Replacement from input-file-mask to produce output file')
	args = parser.parse_args()
	if args.plot_filter:
		s = Signal(args.input_file, window_len=25, frame_rate=100)
	elif args.input_file and args.output_file:
		run_front_end(args.input_file, args.output_file)
	else:
		front_end_mpi(args.input_file_list,
				args.input_file_mask,
			args.replacement)
